# test.py
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'

import cv2 as cv
import numpy as np
import glob

from model import build_encoder_decoder, build_refinement
from data_generator import normalize_input, denormalize_output, depth_random_scale_shift

import h5py as h5
import tensorflow as tf

if __name__ == '__main__':
    img_rows, img_cols = 288, 384
    model_path = '../Models/DIM/coarse.00000079-0.0050.hdf5'
    model = build_encoder_decoder(img_rows, img_cols, train=True)
    #model = build_refinement(model, train=False)
    #model.summary()

    f = h5.File(model_path, 'r')
    w = f['/model_weights']
    params = dict()
    for k in w.keys():
        if len(w[k].keys()) > 0:
            sub_keys = w[k].keys()
            for k2 in sub_keys:
                assert k == k2
                for k3 in w[k][k2].keys():
                    name_ = k + '/' + k3
                    val_ = w[k][k2][k3][()]
                    params[name_] = val_
                    if np.any(np.isnan(val_)):
                        print('NAN VAR FOUND: ')
                        print(name_)
                        #val_[np.where(np.isnan(val_))] = 1.0E-20
                    if np.any(np.isinf(val_)):
                        print('INF VAR FOUND: ')
                        print(name_)
                        #val_[np.where(np.isinf(val_))] = 0

    #print(params)
    vars = tf.trainable_variables()
    assert len(params) == len(vars)

    sess = tf.Session()

    t_deconv6 = sess.graph.get_tensor_by_name('deconv6/Relu:0')

    for i in range(len(vars)):
        sess.run(vars[i].assign(params[vars[i].name]))

    # check all uninitialized variables
    var_unset = tf.report_uninitialized_variables(tf.global_variables())
    print(sess.run(var_unset))
    #fine.load_weights(model_path)
    t_out = model.outputs[0]
    t_in = model.inputs[0]

    files_rgbd = glob.glob('../Datasets/DIM/test/rgbd-1/*.PNG')
    x_test = np.zeros([1, img_rows, img_cols, 3], dtype=np.float32)

    batch_num = len(files_rgbd)

    for i in range(batch_num):
        print(files_rgbd[i])
        rgbd = cv.imread(files_rgbd[i], -1)
        rgbd = cv.resize(rgbd, (img_cols, img_rows))

        rgb = rgbd[:,:,:3]
        x_test[0, :, :, :3] = normalize_input(rgb)
        out, deconv6 = sess.run([t_out, t_deconv6], feed_dict={t_in: x_test})
        assert not np.any(np.isnan(deconv6))
        print(np.max(deconv6))
        #assert not np.any(np.isnan(out))

        out = denormalize_output(out[0,:,:,0])
        out = np.stack([out] * 3, axis=-1)
        #merged = np.concatenate((rgb, out, gt), axis=1)
        merged = np.concatenate((rgb, out), axis=1)
        cv.imwrite(files_rgbd[i].replace('rgbd-1', 'out'), merged)

